# import xarray as xr
# import numpy as np
# from scipy.interpolate import griddata
# import netCDF4 as nc
# import pandas as pd
#
# # 读取FY4A坐标（只需读取一次）
# coord_file = nc.Dataset('/home/liudd/data_preprocessing/FY4A_coordinates.nc', 'r')
# lon_fy4a = coord_file.variables['lon'][:, :].T
# lat_fy4a = coord_file.variables['lat'][:, :].T
# lon_fy4a_360 = lon_fy4a % 360
#
# # 定义通用插值函数
# def interpolate_to_fy4a(data_var, lon_era5, lat_era5):
#     # 生成ERA5网格点
#     lon_grid, lat_grid = np.meshgrid(lon_era5, lat_era5)
#     points = np.column_stack((lon_grid.ravel(), lat_grid.ravel()))
#
#     # 准备目标点
#     target = np.column_stack((lon_fy4a_360.ravel(), lat_fy4a.ravel()))
#     valid = ~np.isnan(target).any(axis=1)
#
#     # 执行插值
#     result = np.full(lon_fy4a.shape, np.nan)
#     interpolated = griddata(
#         points,
#         data_var.values.ravel(),
#         target[valid],
#         method='linear',
#         fill_value=np.nan
#     )
#     result.ravel()[valid] = interpolated
#     return result
#
# # 读取包含多个时间点数据的地表数据文件
# era5_surface = xr.open_dataset('/mnt/datastore/liudddata/ERA5/20200101_Ttest/2020010100_23.nc')
# sp = era5_surface['sp']
# t2m = era5_surface['t2m']
# lon_era5 = era5_surface.longitude.values
# lat_era5 = era5_surface.latitude.values
# times = era5_surface.valid_time.values
#
# # 读取包含多个时间点数据的比湿数据文件
# era5_q = xr.open_dataset('/mnt/datastore/liudddata/ERA5/20200101_Ttest/Specific humidity2020010100_23.nc')
# q = era5_q['q']
# levels = era5_q.pressure_level.values.astype(float)
#
# # 遍历每个时间点
# for i in range(len(times)):
#     print(f"Processing time index {i}...")
#
#     # 提取当前时间点的数据
#     current_sp = sp.isel(valid_time=i)
#     current_t2m = t2m.isel(valid_time=i)
#     current_q = q.isel(valid_time=i)
#
#     # 处理比湿数据
#     # sp_hPa = current_sp / 100  # Pa转hPa
#
#     # def find_nearest_level(p):
#     #     return levels[np.abs(levels - p).argmin()]
#     #
#     # nearest_levels = xr.apply_ufunc(
#     #     find_nearest_level,
#     #     sp_hPa,
#     #     input_core_dims=[[]],
#     #     output_core_dims=[[]],
#     #     vectorize=True,
#     #     dask='parallelized',
#     #     output_dtypes=[levels.dtype]
#     # )
#     surface_q = q.sel(pressure_level=1000, drop=True)
#     # surface_q = current_q.sel(pressure_level=nearest_levels)
#     # min_level = levels.min()
#     # surface_q = surface_q.where(sp_hPa >= min_level, current_q.sel(pressure_level=min_level))
#
#     # 执行插值
#     t2m_interp = interpolate_to_fy4a(current_t2m, lon_era5, lat_era5)
#     sp_interp = interpolate_to_fy4a(current_sp, lon_era5, lat_era5)
#     q_interp = interpolate_to_fy4a(surface_q, lon_era5, lat_era5)
#
#     # 创建当前时间点的数据集
#     ds = xr.Dataset(
#         {
#             "t2m": (("y", "x"), t2m_interp),
#             "sp": (("y", "x"), sp_interp),
#             "q_surface": (("y", "x"), q_interp)
#         },
#         coords={
#             "time": times[i],
#             "lon": (("y", "x"), lon_fy4a),
#             "lat": (("y", "x"), lat_fy4a)
#         }
#     )
#
#     # 添加属性
#     ds['time'].attrs['long_name'] = 'Time'
#     ds['lon'].attrs['units'] = 'degrees_east'
#     ds['lat'].attrs['units'] = 'degrees_north'
#
#     # 生成文件名
#     time_str = pd.to_datetime(times[i]).strftime('%Y%m%d%H')
#     filename = f'/mnt/datastore/liudddata/FY20200101test/ERA5_20200101/ERA5_{time_str}_surface_vars_FY4A_grid.nc'
#
#     # 保存结果
#     ds.to_netcdf(filename)
#     print(f"{time_str}数据处理完成并已保存到 {filename}。")
#
# print("所有逐小时数据处理完成。")

import xarray as xr
import numpy as np
from scipy.interpolate import griddata
import netCDF4 as nc
import pandas as pd

# 读取FY4A坐标（只需读取一次）
coord_file = nc.Dataset('/home/liudd/data_preprocessing/FY4A_coordinates.nc', 'r')
lon_fy4a = coord_file.variables['lon'][:, :].T
lat_fy4a = coord_file.variables['lat'][:, :].T
lon_fy4a_360 = lon_fy4a % 360

# 定义通用插值函数
def interpolate_to_fy4a(data_var, lon_era5, lat_era5):
    # 生成ERA5网格点
    lon_grid, lat_grid = np.meshgrid(lon_era5, lat_era5)
    points = np.column_stack((lon_grid.ravel(), lat_grid.ravel()))

    # 准备目标点
    target = np.column_stack((lon_fy4a_360.ravel(), lat_fy4a.ravel()))
    valid = ~np.isnan(target).any(axis=1)

    # # 检查数据维度
    # print(f"points shape: {points.shape}")
    # print(f"values shape: {data_var.values.ravel().shape}")

    # 执行插值
    result = np.full(lon_fy4a.shape, np.nan)
    interpolated = griddata(
        points,
        data_var.values.ravel(),
        target[valid],
        method='linear',
        fill_value=np.nan
    )
    result.ravel()[valid] = interpolated
    return result

# 读取包含多个时间点数据的地表数据文件
# era5_surface = xr.open_dataset('/mnt/datastore/liudddata/ERA5/20200101_Ttest/2020010100_23.nc')
era5_surface = xr.open_dataset('/mnt/datastore/liudddata/ERA5/20200104/temp20200401-30.nc')
sp = era5_surface['sp']
t2m = era5_surface['t2m']
lon_era5 = era5_surface.longitude.values
lat_era5 = era5_surface.latitude.values
times = era5_surface.valid_time.values

# 读取包含多个时间点数据的比湿数据文件
# era5_q = xr.open_dataset('/mnt/datastore/liudddata/ERA5/20200101_Ttest/Specific humidity2020010100_23.nc')
era5_q = xr.open_dataset('/mnt/datastore/liudddata/ERA5/20200104/specifichumidity_1000hPa20200401_31.nc')
q = era5_q['q']
levels = era5_q.pressure_level.values.astype(float)

# 遍历每个时间点
for i in range(len(times)):
    print(f"Processing time index {i}...")

    # 提取当前时间点的数据
    current_sp = sp.isel(valid_time=i)
    current_t2m = t2m.isel(valid_time=i)
    current_q = q.isel(valid_time=i)

    # 处理比湿数据
    surface_q = q.sel(pressure_level=1000, drop=True).isel(valid_time=i)

    # # 检查 surface_q 的维度
    # print(f"surface_q shape: {surface_q.shape}")

    # 执行插值
    t2m_interp = interpolate_to_fy4a(current_t2m, lon_era5, lat_era5)
    sp_interp = interpolate_to_fy4a(current_sp, lon_era5, lat_era5)
    q_interp = interpolate_to_fy4a(surface_q, lon_era5, lat_era5)

    # 创建当前时间点的数据集
    ds = xr.Dataset(
        {
            "temp_2m": (("y", "x"), t2m_interp),
            "surface_pressure": (("y", "x"), sp_interp),
            "surface_specific_humidity": (("y", "x"), q_interp)
        },
        coords={
            "time": times[i],
            "lon": (("y", "x"), lon_fy4a),
            "lat": (("y", "x"), lat_fy4a)
        }
    )

    # 添加属性
    ds['time'].attrs['long_name'] = 'Time'
    ds['lon'].attrs['units'] = 'degrees_east'
    ds['lat'].attrs['units'] = 'degrees_north'

    # 生成文件名
    time_str = pd.to_datetime(times[i]).strftime('%Y%m%d%H')
    filename = f'/mnt/datastore/liudddata/ERA5_output/202001_04/ERA5_{time_str}_surface_vars_FY4A_grid.nc'

    # 保存结果
    ds.to_netcdf(filename)
    print(f"{time_str}数据处理完成并已保存到 {filename}。")

print("所有逐小时数据处理完成。")